import os
import random
from predict import final_predict_files

def detect_malware(file_path):
    """
    检测给定文件的恶意软件类型。

    :param file_path: 要检测的文件路径
    :return: 预测的恶意软件类型和概率
    """
    class_labels, predicted_probabilities = final_predict_files(file_path)
    
    # 设定阈值
    threshold = 0.70
    
    # 确定是否有类型的概率大于阈值
    detected_type = None
    for label, probability in zip(class_labels, predicted_probabilities):
        if probability > threshold:
            detected_type = label
            break
    
    return detected_type

def find_files_with_extensions(directory, extensions):
    """
    查找指定目录中具有特定扩展名的所有文件。

    :param directory: 要搜索的目录路径
    :param extensions: 需要查找的文件扩展名集合
    :return: 匹配的文件列表
    """
    matched_files = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.endswith(tuple(extensions)):
                matched_files.append(os.path.join(root, file))
    return matched_files

def sample_files_from_directory(directory, num_samples=20):
    """
    从指定目录中随机选择一定数量的文件。

    :param directory: 要选择的目录路径
    :param num_samples: 要选择的文件数量
    :return: 随机选择的文件列表
    """
    extensions = ['.py', '.java', '.go']
    files = find_files_with_extensions(directory, extensions)
    
    if len(files) > num_samples:
        return random.sample(files, num_samples)
    else:
        return files

def calculate_false_positive_rate(root_directory):
    """
    计算给定目录中各类别文件的误报率。

    :param root_directory: 根目录路径
    :return: 误报率
    """
    folder_names = ['normal', 'backdoor', 'dos_attack', 'ftp_exfiltration', 'keylogger', 'persistence_backdoor', 'ransomware', 'reverse_shell']
    total_files = 0
    false_positive_count = 0
    
    for folder in folder_names:
        folder_path = os.path.join(root_directory, folder)
        if os.path.isdir(folder_path):
            sampled_files = sample_files_from_directory(folder_path)
            total_files += len(sampled_files)
            for file in sampled_files:
                detected_type = detect_malware(file)
                if detected_type != folder:
                    false_positive_count += 1

    if total_files == 0:
        return 0.0
    
    return false_positive_count / total_files

def main():
    # 指定要搜索的根目录
    root_directory_path = '../data/train_data'

    # 计算并打印误报率
    false_positive_rate = calculate_false_positive_rate(root_directory_path)

    return false_positive_rate

    # # ANSI 转义码：\033[31m 是红色，\033[0m 是重置
    # print(f"\033[31m系统误报率为：{false_positive_rate:.2%}\033[0m")